import os
from pydub import AudioSegment
import srt
from datetime import timedelta
import re
import numpy as np
from pydub.silence import split_on_silence
import shutil
import atexit
import subprocess

# 定义静默子进程创建参数（Windows）
CREATE_NO_WINDOW = 0x08000000

# Monkey patch subprocess.Popen，所有子进程都静默（包括 pydub 内部的 ffmpeg）
_popen_original = subprocess.Popen


def _popen_silent(*args, **kwargs):
    kwargs["creationflags"] = CREATE_NO_WINDOW
    return _popen_original(*args, **kwargs)

subprocess.Popen = _popen_silent

# 配置 pydub 的 ffmpeg 参数
AudioSegment.converter = "ffmpeg"
AudioSegment.ffmpeg = "ffmpeg"
AudioSegment.ffmpeg_params = ["-hide_banner", "-loglevel", "error"]

# 存储所有创建的临时目录
temp_dirs = set()

def cleanup_temp_dirs():
    """清理所有临时目录"""
    for temp_dir in temp_dirs:
        try:
            if os.path.exists(temp_dir):
                shutil.rmtree(temp_dir)
                print(f"已清理临时目录: {temp_dir}")
        except Exception as e:
            print(f"清理临时目录时出错: {str(e)}")

# 注册退出时的清理函数
atexit.register(cleanup_temp_dirs)

def time_to_milliseconds(time_obj):
    """Convert a timedelta object to milliseconds."""
    return int(time_obj.total_seconds() * 1000)

def get_file_extension(file_path):
    """Get the file extension from the path."""
    return os.path.splitext(file_path)[1].lower()

def convert_srt_time_format(srt_content):
    """Convert SRT time format from comma to dot separator."""
    pattern = r'(\d{2}:\d{2}:\d{2}),(\d{3})'
    return re.sub(pattern, r'\1.\2', srt_content)

def export_segment(segment, segment_path, input_format, output_format):
    """导出单个音频片段的函数"""
    try:
        # 规范化路径，确保使用正斜杠
        segment_path = os.path.normpath(segment_path).replace('\\', '/')
        #print(f"导出路径: {segment_path}")  # 添加调试信息
        
        # 如果输入输出格式相同，直接导出而不重新编码
        if input_format == output_format:
            segment.export(segment_path, format=output_format)
        else:
            # 设置导出参数
            export_params = {}
            if output_format == "mp3":
                export_params = {
                    "format": "mp3",
                    "codec": "libmp3lame",
                    "bitrate": "192k"
                }
            else:  # wav
                export_params = {
                    "format": "wav",
                    "codec": "pcm_s16le"
                }
            
            # 导出音频片段
            segment.export(segment_path, **export_params)
        return True
    except Exception as e:
        print(f"导出片段时出错: {str(e)}")
        return False

def split_audio_by_volume(audio_path, output_dir, min_silence_len=1000, silence_thresh=-30, 
                         progress_callback=None, output_format="wav"):
    """基于音量分割音频"""
    try:
        # 规范化路径
        audio_path = os.path.normpath(audio_path).replace('\\', '/')
        output_dir = os.path.normpath(output_dir).replace('\\', '/')
        
        # Create temp directory if it doesn't exist
        temp_dir = os.path.join(output_dir, ".tmp")
        temp_dir = os.path.normpath(temp_dir).replace('\\', '/')
        os.makedirs(temp_dir, exist_ok=True)
        temp_dirs.add(temp_dir)  # 添加到临时目录集合
        
        # 获取输入文件格式
        input_format = get_file_extension(audio_path)[1:].lower()  # 去掉点号
        
        # Load audio file
        if progress_callback:
            progress_callback("正在加载音频文件...", 0, 1, temp_dir)
        audio = AudioSegment.from_file(audio_path)
        
        if progress_callback:
            progress_callback("正在分析音频...", 0, 1, temp_dir)
            
        # 使用split_on_silence函数分割音频
        chunks = split_on_silence(
            audio,
            min_silence_len=min_silence_len,  # 最小静音长度（毫秒）
            silence_thresh=silence_thresh,    # 静音阈值（dB）
            keep_silence=100,                 # 保留静音片段的长度（毫秒）
            seek_step=25                      # 每次检查的步长（毫秒）
        )
        
        # 导出音频片段
        total_segments = len(chunks)
        if progress_callback:
            progress_callback("正在导出音频片段...", 0, total_segments, temp_dir)
        
        for i, chunk in enumerate(chunks, 1):
            segment_filename = f"{i:03d}.{output_format}"
            segment_path = os.path.join(temp_dir, segment_filename)
            segment_path = os.path.normpath(segment_path).replace('\\', '/')
            export_segment(chunk, segment_path, input_format, output_format)
            if progress_callback:
                progress_callback(f"{i:03d}.{output_format}", i, total_segments, temp_dir)
        
        return temp_dir
            
    except Exception as e:
        print(f"Error in split_audio_by_volume: {str(e)}")
        raise

def split_audio_by_srt(audio_path, srt_path, output_dir, group_size=1, progress_callback=None, output_format="wav"):
    """基于字幕文件分割音频"""
    try:
        # 规范化路径
        audio_path = os.path.normpath(audio_path).replace('\\', '/')
        srt_path = os.path.normpath(srt_path).replace('\\', '/')
        output_dir = os.path.normpath(output_dir).replace('\\', '/')
        
        # Create temp directory if it doesn't exist
        temp_dir = os.path.join(output_dir, ".tmp")
        temp_dir = os.path.normpath(temp_dir).replace('\\', '/')
        os.makedirs(temp_dir, exist_ok=True)
        temp_dirs.add(temp_dir)  # 添加到临时目录集合
        
        # 获取输入文件格式
        input_format = get_file_extension(audio_path)[1:].lower()  # 去掉点号
        
        # Load audio file
        if progress_callback:
            progress_callback("正在加载音频文件...", 0, 1, temp_dir)
        audio = AudioSegment.from_file(audio_path)
        
        # Parse SRT file
        if progress_callback:
            progress_callback("正在解析字幕文件...", 0, 1, temp_dir)
            
        with open(srt_path, 'r', encoding='utf-8') as f:
            srt_content = f.read()
        
        # Convert time format from comma to dot separator
        srt_content = convert_srt_time_format(srt_content)
        
        subtitles = list(srt.parse(srt_content))
        
        # Group subtitles
        grouped_subtitles = []
        for i in range(0, len(subtitles), group_size):
            group = subtitles[i:i + group_size]
            if group:
                # Create a new subtitle with the combined time range
                start_time = group[0].start
                end_time = group[-1].end
                content = "\n".join(sub.content for sub in group)
                grouped_subtitles.append(srt.Subtitle(
                    index=len(grouped_subtitles) + 1,
                    start=start_time,
                    end=end_time,
                    content=content
                ))
        
        # 导出音频片段
        total_segments = len(grouped_subtitles)
        if progress_callback:
            progress_callback("正在导出音频片段...", 0, total_segments, temp_dir)
        
        for i, subtitle in enumerate(grouped_subtitles, 1):
            start_time = time_to_milliseconds(subtitle.start)
            end_time = time_to_milliseconds(subtitle.end)
            
            # Extract the audio segment
            segment = audio[start_time:end_time]
            
            # 根据指定格式导出
            segment_filename = f"{i:03d}.{output_format}"
            segment_path = os.path.join(temp_dir, segment_filename)
            segment_path = os.path.normpath(segment_path).replace('\\', '/')
            export_segment(segment, segment_path, input_format, output_format)
            
            if progress_callback:
                progress_callback(f"{i:03d}.{output_format}", i, total_segments, temp_dir)
        
        return temp_dir
            
    except Exception as e:
        print(f"Error in split_audio_by_srt: {str(e)}")
        raise 